import numpy as np

def search_documents(query, index, documents, embedding_model, top_k=2):
    """
    根据用户查询在向量数据库中检索最相关的文档片段。

    Args:
        query: 用户查询 (字符串)
        index: FAISS 向量数据库索引
        documents: 文档列表 (与索引对应的原始文档)
        embedding_model: 嵌入模型
        top_k: 返回最相关的文档数量

    Returns:
        list: 最相关的文档片段列表
    """
    query_embedding = embedding_model.encode([query]) # 将用户查询转换为嵌入向量
    D, I = index.search(np.array(query_embedding), top_k) # 在向量数据库中搜索最相似的 top_k 个文档

    # I 是索引，D 是距离 (相似度)，按距离排序
    relevant_documents = [documents[i] for i in I[0]] # 根据索引获取对应的原始文档

    return relevant_documents

if __name__ == '__main__':
    from sentence_transformers import SentenceTransformer
    import faiss
    from data_indexing import create_index_and_embeddings # 假设 data_indexing.py 在同一目录下

    # 模拟企业文档 (需要和 data_indexing.py 中保持一致，或者从文件/数据库加载)
    documents = [
        "我们的退货政策是客户可以在30天内退货。",
        "2025年新的假期安排，每人享有法定年假12天。",
        "上个月，我们的销售政策是全场商品8折优惠。",
        "公司鼓励员工参与志愿活动，每年提供一天带薪志愿假。"
    ]

    # 加载嵌入模型 (需要和 data_indexing.py 中使用的模型一致)
    embedding_model = SentenceTransformer('all-MiniLM-L12-v2')

    # 创建索引 (或者从文件加载，如果之前保存了)
    index, _ = create_index_and_embeddings(documents) # 只需要 index，忽略 embeddings

    # 测试查询
    query = "销售政策"
    relevant_docs = search_documents(query, index, documents, embedding_model)

    print(f"查询: {query}")
    print("相关文档:")
    for doc in relevant_docs:
        print(f"- {doc}")